#include <traceUtils/outputUtils.idc>
#include <traceUtils/forRunToAddr.idc>
/**
 * 主要属性说明：
 * this.canDo            代表是否初始化完成，值为0或者1，为0代表初始化失败，为1代表初始化成功
 * this.real_addr_fp     输出主文件真实地址的句柄， 加偏移后的地址
 * this.orig_addr_fp     输出主文件原始地址的句柄， 不加偏移后的地址
 * this.text_min_addr    主程序text段开始地址
 * this.text_max_addr    主程序text段结束地址
 * this.ASL_addr         偏移地址
 * this.blackFuncNames   函数黑名单数组，在函数黑名单里面的函数在识别到B或者BL指令时会校验调换函数，匹配黑名单使用StepOver跳过
 * this.sp               输出汇编代码指令最后行数
 * this.outputUtils      输出工具类 
*/
class TraceHook{
    /**
     * stopAddr               trace end address (no offset address)
     * sleepTime              step sleep time (s)
     * moduleName             execution binary file name and output file name
     * baseOuputFolderPath    base output folder path (folder must exist)
     */
    TraceHook(stopAddr, sleepTime, moduleName, baseOuputFolderPath){
        this.stopAddr = stopAddr;
        this.sleepTime = sleepTime;
        this.moduleName = moduleName;
        this.baseOuputFolderPath = baseOuputFolderPath;

        Message("TraceHook init start\n");
        this.canDo = 1;
        this.sp = 0;
        if(this.determineTextSegment() != 1){
            Warning("No offset found ！！！exit ！！！");
            this.canDo = 0;
            return;
        }

        this.forRunToAddr = ForRunToAddr();
        this.determineLibSystemPthreadSegment();

        this.systemSkip();

        this.blackFunc();

        this.objcMsgSendBlackFunc = Array(AR_STR, "objcMsgSendBlack");
        Message("TraceHook init end\n");
    }

    ~TraceHook(){
        Message("get trace end!\n");
    }

    trace(){
        auto code, addr, mnem, inst, disAsm;
        EnableTracing(TRACE_STEP, 1);
        code = GetDebuggerEvent(WFNE_ANY | WFNE_CONT, -1);
        for( ;code > 0; code = GetDebuggerEvent(WFNE_ANY | WFNE_CONT, -1)){
            addr = GetEventEa();
            if(addr == this.stopAddr + this.text_min_addr){
                break;
            }

            if(addr < this.text_min_addr || addr > this.text_max_addr){
                continue;
            }

            this.sp = this.sp + 1;
            this.outputUtils.outputByAddr(addr, this.sp);
        }
        EnableTracing(TRACE_STEP, 0);
    }

    pdObjcMsgSendBlack(disAsm){
        if(strstr(disAsm, "_objc_msgSend") == -1){
            return 0;
        }

        this.outputUtils.registerValue.addrType = 1;
        auto realResult = this.outputUtils.registerValue.getValue("X1", 80);
        if(this.objcMsgSendBlackFunc.contain(realResult) != -1){
            Message("filter _objc_msgSend X1=%s \n", realResult);
            return 1;
        }
        return 0;
    }

    addMsgSendBlackFunc(name){
        this.objcMsgSendBlackFunc.add(name);
    }

    goBeyondAddrHandle(addr){
        auto segmentName, forEndAddr;
        auto forRunToAddrIndex = this.forRunToAddr.getForStartAddrIndex(addr);
        if(forRunToAddrIndex != -1){
            segmentName = this.forRunToAddr.getSegmentNameByIndex(forRunToAddrIndex);
            forEndAddr = this.forRunToAddr.getForEndAddrByIndex(forRunToAddrIndex);
            Message("%s for handle now addr 0x%x run to 0x%x \n", segmentName, addr, forEndAddr);
            Sleep(this.sleepTime * 1000);
            RunTo(forEndAddr);
            return;
        }

        if(addr < this.libSystemPthreadStart || addr > this.libSystemPthreadEnd){
            Sleep(this.sleepTime * 1000);
            StepOver();
            return;
        }

        Sleep(this.sleepTime * 1000);
        StepInto();
    }

    addForRunToAddr(segmentName, baseAddr, startAddr, endAddr){
        auto baseStartAddr, baseEndAddr;
        baseStartAddr = startAddr - baseAddr;
        baseEndAddr = endAddr - baseAddr;

        auto minAddr, maxAddr;
        minAddr = this.determineSegmentMinAddr(segmentName);
        if(minAddr == 0xffffffffffffffff){
            Warning("add for run to addr get segment min file ！！！ \n");
            return 0;
        }

        auto tmpMinAddr, tmpMaxAddr;
        tmpMinAddr = minAddr + baseStartAddr;
        tmpMaxAddr = minAddr + baseEndAddr;
        this.forRunToAddr.add(segmentName, tmpMinAddr, tmpMaxAddr);
        // this.forRunToAddr.forEcho();
    }

    blackFunc(){
        auto addr, name, segName;
        addr = 0;
        this.blackFuncNames = Array(AR_STR, "blackFuncNames");
        for(addr = NextFunction(addr); addr != BADADDR; addr = NextFunction(addr)){
            name = Name(addr);
            segName = SegName(addr);
            if(segName != "__text" && name != "_objc_msgSend"){
                this.blackFuncNames.add(name);
            }
        }
        //this.blackFuncNames.arrayEcho();
    }

    addBlackFunc(name){
        this.blackFuncNames.add(name);
    }

    determineSegmentMinAddr(moduleName){
        auto x, name;
        for (x = GetFirstModule(); x != BADADDR; x = GetNextModule(x)){
            name = GetModuleName(x);
            if(strstr(name, moduleName) != -1){
                return x;
            }
        }
        return 0xffffffffffffffff;
    }

    determineSegmentMaxAddr(moduleName){
        auto x, name;
        for (x = GetFirstModule(); x != BADADDR; x = GetNextModule(x)){
            name = GetModuleName(x);
            if(strstr(name, moduleName) != -1){
                return x + GetModuleSize(x);
            }
        }
        return 0xffffffffffffffff;
    }

    determineTextSegment(){
        auto minAddr, maxAddr;
        minAddr = this.determineSegmentMinAddr(this.moduleName);
        if(minAddr == 0xffffffffffffffff){
            return 0;
        }

        maxAddr = this.determineSegmentMaxAddr(this.moduleName);
        if(maxAddr == 0xffffffffffffffff){
            return 0;
        }

        this.text_min_addr = minAddr;
        this.text_max_addr = maxAddr;
        this.ASL_addr = minAddr - 0x100000000;
        this.outputUtils = OutputUtils(this.ASL_addr, this.baseOuputFolderPath, this.moduleName);
        if(this.outputUtils.canDo != 1){
            return 0;
        }

        auto logContent = form("[%s] asl addr [0x%x], text_min_addr [0x%x], text_max_addr [0x%x], real stop addr [0x%x]\n", this.moduleName, this.ASL_addr, this.text_min_addr, this.text_max_addr, this.stopAddr + this.text_min_addr);
        Message(logContent);
        this.outputUtils.writeToFile(logContent);
        return 1;
    }

    determineLibSystemPthreadSegment(){
        auto minAddr, maxAddr;
        minAddr = this.determineSegmentMinAddr("libsystem_pthread.dylib");
        if(minAddr == 0xffffffffffffffff){
            return 0;
        }

        maxAddr = this.determineSegmentMaxAddr("libsystem_pthread.dylib");
        if(maxAddr == 0xffffffffffffffff){
            return 0;
        }

        this.libSystemPthreadStart = minAddr;
        this.libSystemPthreadEnd = maxAddr;
        return 1;
    }

    systemSkip(){
        // ---------------- 针对系统类库for循环跳过 --------------------------------
        // this.addForRunToAddr("libsystem_pthread.dylib", 0x0, 0x16D8, 0x16EC);
        // this.addForRunToAddr("libsystem_pthread.dylib", 0x0, 0x16D8, 0x19C4);
        // this.addForRunToAddr("libobjc.A.dylib", 0x0, 0x6164, 0x6174);
        // this.addForRunToAddr("libsystem_malloc.dylib", 0x0, 0x13C8, 0x13D8);
        // this.addForRunToAddr("libsystem_malloc.dylib", 0x0, 0x1AFC, 0x1B10);
        // this.addForRunToAddr("libobjc.A.dylib", 0x0, 0x6268, 0x6294);
        // this.addForRunToAddr("libsystem_pthread.dylib", 0x0, 0x19CC, 0x19E4);
        // this.addForRunToAddr("libsystem_pthread.dylib", 0x0, 0x19F4, 0x1A0C);
        // this.addForRunToAddr("libsystem_pthread.dylib", 0x0, 0x19B0, 0x19C4);
        // this.addForRunToAddr("libobjc.A.dylib", 0x0, 0x200AC, 0x200C4);
        // this.addForRunToAddr("CoreFoundation", 0x0, 0xDF2F8, 0xDF30C);
        // this.addForRunToAddr("libdispatch.dylib", 0x0, 0xD8C8, 0xD8D8);
        this.addForRunToAddr("libobjc.A.dylib", 0x0, 0x241AC, 0x241C4);
        // this.addForRunToAddr("libobjc.A.dylib", 0x18DE04000, 0x18DE281B8, 0x18DE281C4);
        // this.addForRunToAddr("libobjc.A.dylib", 0x1845C8000, 0x1845E0614, 0x1845E0624);
        // this.addForRunToAddr("libobjc.A.dylib", 0x1845C8000, 0x1845E024C, 0x1845E0288);
        // this.addForRunToAddr("libobjc.A.dylib", 0x1845C8000, 0x1845E0294, 0x1845E02B8);
        // this.addForRunToAddr("libobjc.A.dylib", 0x1845C8000, 0x1845DFED8, 0x1845DFF20);
        // this.addForRunToAddr("libobjc.A.dylib", 0x1845C8000, 0x1845DFF30, 0x1845DFF44);
        // this.addForRunToAddr("libobjc.A.dylib", 0x1845C8000, 0x1845E971C, 0x1845E9788);
        // this.addForRunToAddr("libobjc.A.dylib", 0x1845C8000, 0x1845E00F4, 0x1845E0104);
        // this.addForRunToAddr("libobjc.A.dylib", 0x1845C8000, 0x1845E0150, 0x1845E0164);
        // this.addForRunToAddr("libobjc.A.dylib", 0x1845C8000, 0x1845EAD5C, 0x1845EAD74);
        // this.addForRunToAddr("libsystem_pthread.dylib", 0x184C2A000, 0x184C2B3D0, 0x184C2B3E4);
        // this.addForRunToAddr("CoreFoundation", 0x185A69000, 0x185B48490, 0x185B484A4);
        // this.addForRunToAddr("CoreFoundation", 0x185A69000, 0x185B48330, 0x185B48344);
        // this.addForRunToAddr("libsystem_platform.dylib", 0x184C23000, 0x184C256A8, 0x184C256C0);
        // this.addForRunToAddr("libsystem_platform.dylib", 0x184C23000, 0x184C24C48, 0x184C24C58);
        this.addForRunToAddr("CoreFoundation", 0x0, 0xDCBC8, 0xDCBDC);
    }
};